package mmodel

import (
	"sync"
)

type Set[T string | int64 | int32 | int16 | int8 | int | float64 | float32 | uint | uint8 | uint16 | uint32 | uint64 | uintptr | complex64 | complex128 | bool] map[T]bool

func (ce Set[T]) AddAll(keys ...T) {
	for i := range keys {
		(ce)[keys[i]] = true
	}
}
func (ce Set[T]) Add(key T) {
	(ce)[key] = true
}
func (ce Set[T]) Remove(key T) {
	delete(ce, key)
}
func (ce Set[T]) RemoveAll(keys ...T) {
	for i := range keys {
		delete(ce, keys[i])
	}
}
func (ce Set[T]) Val() []T {
	res := make([]T, 0)
	for k, _ := range ce {
		res = append(res, k)
	}
	return res
}
func (ce Set[T]) Contains(key T) bool {
	return (ce)[key]
}

func NewSyncSet[T string | int64 | int32 | int16 | int8 | int | float64 | float32 | uint | uint8 | uint16 | uint32 | uint64 | uintptr | complex64 | complex128 | bool]() *ConcurrentSet[T] {
	return &ConcurrentSet[T]{
		lock:  &sync.RWMutex{},
		locks: map[T]*sync.RWMutex{},
		Data:  &Set[T]{},
	}
}

type ConcurrentSet[T string | int64 | int32 | int16 | int8 | int | float64 | float32 | uint | uint8 | uint16 | uint32 | uint64 | uintptr | complex64 | complex128 | bool] struct {
	Data  *Set[T]
	lock  *sync.RWMutex
	locks map[T]*sync.RWMutex
}

func (ce *ConcurrentSet[T]) AddAll(keys ...T) {
	for i := range keys {
		ce.Add(keys[i])
	}
}
func (ce *ConcurrentSet[T]) Add(key T) {
	lock := ce.getLock(key)
	lock.Lock()
	(*ce.Data)[key] = true
	lock.Unlock()
}
func (ce *ConcurrentSet[T]) Remove(key T) {
	lock := ce.getLock(key)
	lock.Lock()
	delete(*ce.Data, key)
	lock.Unlock()
}
func (ce *ConcurrentSet[T]) RemoveAll(keys ...T) {
	for i := range keys {
		ce.Remove(keys[i])
	}
}
func (ce *ConcurrentSet[T]) Val() []T {
	res := make([]T, 0)
	for k, _ := range *ce.Data {
		res = append(res, k)
	}
	return res
}

func (ce *ConcurrentSet[T]) Map() map[T]bool {
	return *ce.Data
}
func (ce *ConcurrentSet[T]) getLock(key T) *sync.RWMutex {
	var result *sync.RWMutex
	ce.lock.Lock()
	result = ce.locks[key]
	if result == nil {
		result = &sync.RWMutex{}
		ce.locks[key] = result
	}
	ce.lock.Unlock()
	return result
}
